#coding=utf-8

from __future__ import division
from __future__ import absolute_import
from __future__ import print_function

import os, sys
from PIL import Image
import numpy as np
import torch
import onnx
import onnxruntime
import cv2

class ONNXModel():
    def __init__(self, onnx_path):
        self.onnx_session = onnxruntime.InferenceSession(onnx_path)
        self.input_name = self.get_input_name(self.onnx_session)
        self.output_name = self.get_output_name(self.onnx_session)
    def get_output_name(self, onnx_session):
        output_name = []
        for node in onnx_session.get_outputs():
            output_name.append(node.name)
        return output_name
    def get_input_name(self, onnx_session):
        input_name = []
        for node in onnx_session.get_inputs():
            input_name.append(node.name)
        return input_name
    def get_input_feed(self, input_name, image_numpy):
        input_feed = {}
        for name in input_name:
            input_feed[name] = image_numpy
        return input_feed
    def forward(self, image_numpy):
        input_feed = self.get_input_feed(self.input_name, image_numpy)
        segmap = self.onnx_session.run(self.output_name, input_feed=input_feed)
        return segmap
def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

if __name__ == '__main__':


    model_file = '/data/deling/DSFD/FaceDetection-DSFD_Full/dsfd.onnx'  # '/home/adapt/efficientnet/efficientnet-b0.onnx'
    prep_file_path = '/data/deling/DSFD/FaceDetection-DSFD_Full/onnx_images/' # '/home/adapt/efficientnet/prep_dataset'
    output_path = '/data/deling/DSFD/FaceDetection-DSFD_Full/onnx_out/'  # '/home/adapt/efficientnet/infer_onnx_res'

    if not os.path.exists(output_path):
        os.makedirs(output_path)
    net = ONNXModel(model_file)
    files = os.listdir(prep_file_path)

    # img_ori = cv2.imread('/data/deling/DSFD/FaceDetection-DSFD_Full/onnx_images/0_Parade_marchingband_1_20.jpg')
    # img = cv2.resize(img_ori, (224, 224), interpolation=cv2.INTER_CUBIC)
    # img_input = img[..., ::-1]  # BGR to RGB
    # img_input = (np.float32(img) / 255.0 - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
    # img_input = img_input.transpose((2, 0, 1))
    # img_input = torch.from_numpy(img_input).unsqueeze(0)
    # img_input = img_input.type(torch.FloatTensor)
    #
    #
    # print("##"*10)
    # out = net.forward(to_numpy(img_input))
    # print("**"*10)
    # print(out)

    # img = np.fromfile(os.path.join(prep_file_path, file), dtype='float32')
    # img  = np.reshape(img, (3, 224, 224))

    img_ori = cv2.imread('/data/deling/DSFD/FaceDetection-DSFD_Full/onnx_images/0_Parade_marchingband_1_20.jpg')
    img = cv2.resize(img_ori, (224, 224), interpolation=cv2.INTER_CUBIC)
    img = torch.tensor(img, dtype = torch.float32)
    img = img.unsqueeze(0)
    print("get img shape", img.shape)
    img = img.transpose(1,3)

    #img = img.view(1,3,224,224)
    print("get img shape", img.shape)
    output = net.forward(to_numpy(img))[0]
    output = np.array(output)
    print("##" * 10)
    print(output)
    #np.savetxt(os.path.join(output_path, file.split('.')[0] + '_1.txt'), output, fmt='%.6f')

